#ifndef RSA_H_
#define RSA_H_
#include <iostream>
#include <vector>
#include <string>
#include <gmp.h>
#include <gmpxx.h>
#include "prime.hpp"
#include "math.hpp"

using namespace std;
/*
 * 函数说明
 * RSA::random_prime: 生成随机素数
 * RSA::RSA 构造函数需要传入两个素数
 * RSA::encrypt_num 加密数字(数字需要小于两个素数的乘积)
 * RSA::decrypt_num 解密数字
 *
 * 算法详解：
 * n = p * q, p和q是素数, phi_n是n的欧拉函数, 即(m^phi_n)%n=1
 * e是小于phi_n的素数, d是e模phi_n的乘法逆元，即(e*d)%phi_n = 1, 即e*d=h*phi_n+1
 * m是明文，加密后的密文为c=(m^e)%n.
 * 解密过程 decrypt_msg=c^d%n = m^(e*d)%n=m(h*phi_n+1)%n=m%n
 * m必须要小于n 
 */
class RSA{
private:
	Prime prime;
	mpz_class n, p, q, phi_n, e, d;

	int init_e(){
		vector<int> primes = prime.get_primes();
		for(int i = primes.size() - 1; i >= 0; i--){
			e = primes[i];
			if(phi_n % e != 0 && e < phi_n){
				break;
			}
		}
	}
public:
	RSA(mpz_class _p, mpz_class _q){
		//p = "100001623";
		//q = "100001651";
		p = _p;
		q = _q;
		n = p * q;
		phi_n = (p - 1) * (q - 1);
		mpz_class tmp;
		init_e();
		//这里如果e和phi_n不互质则会有问题
		ext_gcd(e, phi_n, d, tmp);
		//d必须要为正数
		if(d < 0){
			mpz_class f = abs(d / phi_n);
			if(d % phi_n != 0){
				f += 1;
			}
			d += (f * phi_n);
		}
	}

	// 加密数字
	mpz_class encrypt_num(mpz_class m){
		return Math::power_mod(m, e, n);
	}

	// 解密数字
	mpz_class decrypt_num(mpz_class c){
		return Math::power_mod(c, d, n);
	}

	/*
	 * 随机素数，返回0表是没有找到这样的素数
	 */
	static mpz_class random_prime(int bit_len){
		mpz_class zero(0), p;
		int max_try_time = 10000;
		Prime prime;
		for(int i = 0; i < max_try_time; i++){
			p = prime.pseudo_random_prime(bit_len);
			if(p != zero){
				return p;
				// RSA验证素数
				if(check_is_prime(p)){
					return p;
				}
			}
		}
		return zero;
	}
	
	//RSA算法验证p是否是素数
	static bool check_is_prime(mpz_class p){
		mpz_class q(1000000007), plaintext(100);
		RSA rsa(p, q);
		mpz_class cipher = rsa.encrypt_num(plaintext);
		return rsa.decrypt_num(cipher) == plaintext;
	}
	
	// 展欧几里得求m和n的最大公约数g=a*x+b*y,特别地，如果a和b互质，则x是a模b的乘法逆元
	static mpz_class ext_gcd(mpz_class a, mpz_class b, mpz_class &x, mpz_class &y){
		if(a % b == 0){
			x = 1;
			y = 1 - a/b;
			return b;
		}else if(b % a == 0){
			x = 1 - b / a;
			y = 1;
			return a;
		}else{
			mpz_class x1, y1;
			mpz_class g = ext_gcd(b, a % b, x1, y1);
			x = y1;
			y = x1 - (a/b) * y1;
			return g;
		}
	}

	//批量测试RSA
	static void batch_test(){
		int max_times = 10;
		int success_num = 0, fail_num = 0;
		for(int i = 0; i < max_times; i++){
			mpz_class p = RSA::random_prime(1024);
			mpz_class q = RSA::random_prime(1024);
			RSA rsa(p, q);
			mpz_class m = Math::random_num(100);
			mpz_class c = rsa.encrypt_num(m);
			if(rsa.decrypt_num(c) == m){
				success_num++;
			}else{
				fail_num++;
			}
			cout << "成功数:" << success_num << " 失败数:" << fail_num << endl;
		}
	}

};

#endif
